from typing import List, Optional, Tuple, Any, Iterable
from langgraph.store.base import BaseStore, Item, Op, Result
from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType, utility
from redis import Redis
import json
from datetime import datetime
import uuid
from langgraph_memory.src.memory.embedding import EmbeddingModel
from langgraph_memory.src.config.config import Settings

class RedisMilvusStore(BaseStore):
    async def abatch(self, ops: Iterable[Op]) -> list[Result]:
        pass

    def batch(self, ops: Iterable[Op]) -> list[Result]:
        pass

    def __init__(self, redis_client: Redis, milvus_collection: Collection, embedder: EmbeddingModel):
        self.redis = redis_client
        self.milvus = milvus_collection
        self.embedder = embedder

    def put(self, namespace: Tuple[str, ...], key: str, value: dict) -> None:
        """存储记忆到 Redis，并插入向量到 Milvus"""
        # 1. 序列化存入 Redis
        redis_key = f"memory:{':'.join(namespace)}:{key}"
        self.redis.set(redis_key, json.dumps(value))

        # 2. 生成 embedding 并插入 Milvus
        # 获取value的文本值
        text = value.get("memory", "")
        vector = self.embedder.embed(text)
        # uuid user_id 文本数据 向量
        self.milvus.insert([[key], [namespace[0]], [text], [vector]])

    def search(self, namespace: Tuple[str, ...], query: str, limit: int = 3) -> List[Item]:
        """语义搜索：用 Milvus 找 top-k，再从 Redis 取完整数据"""
        user_id = namespace[0]
        vector = self.embedder.embed(query)

        # Milvus 向量搜索
        results = self.milvus.search(
            data=[vector],
            anns_field="vector",
            param={"metric_type": "L2", "params": {"nprobe": 10}},
            limit=limit,
            expr=f"user_id == '{user_id}'",  # 限定用户
            output_fields=["id", "text"]
        )

        items = []
        for hit in results[0]:
            memory_id = hit.entity.get("id")
            redis_key = f"memory:{':'.join(namespace)}:{memory_id}"
            raw = self.redis.get(redis_key)
            if raw:
                # 如果原始数据中包含时间信息，则使用它
                value = json.loads(raw)
                created_time = value.get('created_at', datetime.now().isoformat())
                updated_time = value.get('updated_at', datetime.now().isoformat())

                items.append(Item(
                    key=memory_id,
                    namespace=list(namespace),
                    value=value,
                    created_at=created_time,  # 可扩展为真实时间戳
                    updated_at=updated_time
                ))
        return items

    def get(self, namespace: Tuple[str, ...], key: str) -> Optional[Item]:
        redis_key = f"memory:{':'.join(namespace)}:{key}"
        raw = self.redis.get(redis_key)
        if raw:
            # 如果原始数据中包含时间信息，则使用它
            value = json.loads(raw)
            created_time = value.get('created_at', datetime.now().isoformat())
            updated_time = value.get('updated_at', datetime.now().isoformat())
            return Item(
                key=key,
                namespace=list(namespace),
                value=value,
                created_at=created_time,
                updated_at=updated_time
            )
        return None

    def delete(self, namespace: Tuple[str, ...], key: str) -> None:
        redis_key = f"memory:{':'.join(namespace)}:{key}"
        self.redis.delete(redis_key)
        # Milvus 删除逻辑（可选，通常软删除或定期清理）

# 初始化 Milvus Collection（首次运行时调用）
def init_milvus_collection(embedder: EmbeddingModel):
    connections.connect("default", host=Settings.MILVUS_HOST, port=Settings.MILVUS_PORT)

    collection_name = "memories"
    if not utility.has_collection(collection_name):
        fields = [
            FieldSchema(name="id", dtype=DataType.VARCHAR, is_primary=True, max_length=64),
            FieldSchema(name="user_id", dtype=DataType.VARCHAR, max_length=64),
            FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=512),
            FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=embedder.dim)
        ]
        schema = CollectionSchema(fields, "User memories with embeddings")
        collection = Collection(collection_name, schema)
        index_params = {
            "index_type": "IVF_FLAT",
            "metric_type": "L2",
            "params": {"nlist": 128},
        }
        collection.create_index("vector", index_params)
        collection.load()
        return collection
    else:
        collection = Collection(collection_name)
        collection.load()
        return collection